package com.yoshino.leetcode.interview150.heap;

import java.util.*;

class Solution {

    Random random = new Random();
    public int findKthLargest(int[] nums, int k) {
        return quickSelect(nums, 0, nums.length - 1, k - 1);
    }

    private int quickSelect(int[] nums, int l, int r, int index) {
        // 第 index 个最大的
        int q = randomPartiton(nums, l, r);
        if (q == index) {
            return nums[q];
        } else {
            return q < index ? quickSelect(nums, q + 1, r, index) : quickSelect(nums, l, q - 1, index);
        }
    }

    // 中标
    private int randomPartiton(int[] nums, int l, int r) {
        int i = random.nextInt(r - l + 1) + l;
        swap(nums, i,r);
        return partition(nums, l, r);
    }

    // 快排，中位数为 nums[r]
    private int partition(int[] nums, int l, int r) {
        // 比 x 大的数的个数
        int x = nums[r], i = l - 1;
        for (int j = l; j < r; j++) {
            if (nums[j] > x) {
                // 放在前面
                swap(nums, ++i, j);
            }
        }
        // 交换游标
        swap(nums, i + 1, r);
        // 游标的位置
        return i + 1;
    }

    private void swap(int[] nums, int i, int j) {
        // 交换 i,j
        int temp = nums[i];
        nums[i] = nums[j];
        nums[j] = temp;
    }

    public int findMaximizedCapital(int k, int w, int[] profits, int[] capital) {
        int n = profits.length;
        int cur = 0;
        int[][] arr = new int[n][2];
        for (int i = 0; i < n; i++) {
            arr[i][0] = capital[i];
            arr[i][1] = profits[i];
        }
        Arrays.sort(arr, Comparator.comparingInt(a -> a[0]));
        PriorityQueue<Integer> queue = new PriorityQueue<>((x, y) -> y - x);
        for (int i = 0; i < k; i++) {
            // 当前够支付的情况下的纯利润列表，最大堆
            while (cur < n && arr[cur][0] <= w) {
                queue.add(arr[cur][1]);
                cur++;
            }
            if (!queue.isEmpty()) {
                w += queue.poll();
            } else {
                break;
            }
        }
        return w;
    }

    public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        PriorityQueue<int[]> queue = new PriorityQueue<>(k, (a, b) -> nums1[a[0]] + nums2[a[1]] - nums1[b[0]] - nums2[b[1]]);
        List<List<Integer>> res = new ArrayList<>();
        int m = nums1.length, n = nums2.length;
        // 存入基础的
        for (int i = 0; i < Math.min(m, k); i++) {
            queue.offer(new int[]{i, 0});
        }
        while (k-- > 0 && !queue.isEmpty()) {
            int[] pair = queue.poll();
            List<Integer> temp = new ArrayList<Integer>(){{
                add(nums1[pair[0]]);
                add(nums2[pair[1]]);
            }};
            res.add(temp);
            if (pair[1] + 1 < n) {
                // 拿一个放一个
                queue.add(new int[] {pair[0], pair[1] + 1});
            }
        }
        return res;
    }
}